import time
from datetime import datetime
import numpy as np
import torch
from .callback_manager import Callback, RunContext

from ..optimizer import Updater

class RunInfo(Callback):
    r"""Callback to print MD simulation information

    Args:
        per_steps (int): How many steps to print information once
        per_epoch (int): How many epochs to print information once
        show_total_time (bool): Whether to show total running time
        show_single_time (bool): Whether to show single step running time
        check_force (bool): Whether to check force

    Supported Platforms:
        ``CPU`` ``GPU``
    """

    def __init__(self,
                 per_steps: int = 1,
                 show_total_time: bool = True,
                 show_single_time: bool = False,
                 check_force: bool = False,
                 **kwargs
                 ):
        super().__init__()

        if 'print_freq' in kwargs:
            print("[WARNING] `print_freq` will be removed in future versions, please use "
                  "`per_steps` or `per_epoch` instead")
            per_steps = kwargs['print_freq']

        if not isinstance(per_steps, int) or per_steps < 0:
            raise ValueError(f"per_steps must be a non-negative integer, but got: {per_steps}")

        self.per_steps = per_steps

        self.volume = None
        self.metrics = None

        self.use_pbc = False
        self.use_updater = False

        self.crd = None

        self.begin_time = datetime.now()
        self.step_begin_time = datetime.now()
        self.epoch_begin_time = datetime.now()
        self.print_time = time.time()

        self.show_total_time = show_total_time
        self.show_single_time = show_single_time
        self.check_force = check_force

    def __enter__(self):
        """Return the enter target."""
        return self

    def __exit__(self, *err):
        """Release resources here if have any."""

    def begin(self, run_context: RunContext):
        """
        Called once before the network executing.

        Args:
            run_context (RunContext): Include some information of the model.
        """
        self.begin_time = datetime.now()
        self.print_time = time.time()

        if self.show_total_time:
            print('[TorchSPONGE] Started simulation at', self.begin_time.strftime('%Y-%m-%d %H:%M:%S'))

        self.use_pbc = run_context.pbc_box is not None
        if isinstance(run_context.optimizer, Updater):
            self.use_updater = True

    def epoch_begin(self, run_context: RunContext):
        """
        Called before each epoch beginning.

        Args:
            run_context (RunContext): Include some information of the model.
        """
        self.epoch_begin_time = datetime.now()

    def epoch_end(self, run_context: RunContext):
        """
        Called after each epoch finished.

        Args:
            run_context (RunContext): Include some information of the model.
        """
        pass

    def step_begin(self, run_context: RunContext):
        """
        Called before each step beginning.

        Args:
            run_context (RunContext): Include some information of the model.
        """
        self.step_begin_time = datetime.now()

    def step_end(self, run_context: RunContext):
        """
        Called after each step finished.

        Args:
            run_context (RunContext): Include some information of the model.
        """
        if self.per_steps > 0:
            if run_context.cur_step % self.per_steps == 0:
                self.call_end(run_context)

    def end(self, run_context: RunContext):
        """
        Called once after network training.

        Args:
            run_context (RunContext): Include some information of the model.
        """
        end_time = datetime.now()
        if self.show_total_time:
            print('[TorchSPONGE] Finished simulation at', end_time.strftime('%Y-%m-%d %H:%M:%S'))
            used_time = end_time - self.begin_time
            d = used_time.days
            s = used_time.seconds
            m, s = divmod(s, 60)
            h, m = divmod(m, 60)
            if d >= 1:
                print(f'[TorchSPONGE] Total simulation time: {d:d} days, {h:d} hours, {m:d} minutes and {s:d} seconds.')
            elif h >= 1:
                print(f'[TorchSPONGE] Total simulation time: {h:d} hours, {m:d} minutes and {s:d} seconds.')
            elif m >= 1:
                print(f'[TorchSPONGE] Total simulation time: {m:d} minutes and {s:d} seconds.')
            else:
                print(f'[TorchSPONGE] Total simulation time: {s:d} seconds.')

    def call_begin(self, run_context: RunContext):
        """
        Print information at the beginning.

        Args:
            run_context (RunContext): Include some information of the model.
        """
        print(f'[TorchSPONGE] Step: {run_context.cur_step}')

    def call_end(self, run_context: RunContext):
        """
        Print information at the end.

        Args:
            run_context (RunContext): Include some information of the model.
        """
        if self.show_single_time:
            cur_time = time.time()
            time_str = f'Simulation Time: {cur_time - self.print_time:.3f}s'
            self.print_time = cur_time
        else:
            time_str = ''

        if self.use_updater:
            pe = run_context.potential
            ke = run_context.kinetics
            te = pe + ke

            if self.use_pbc:
                print(f'[TorchSPONGE] Step: {run_context.cur_step}, '
                      f'Time: {run_context.cur_time:.3f}, '
                      f'Total Energy: {self._format_value(te)}, '
                      f'Potential Energy: {self._format_value(pe)}, '
                      f'Kinetic Energy: {self._format_value(ke)}, '
                      f'Temperature: {self._format_value(run_context.temperature)}, '
                      f'Pressure: {self._format_value(run_context.pressure)}, '
                      f'Volume: {self._format_value(run_context.volume)}, ', end=''
                      )
            else:
                print(f'[TorchSPONGE] Step: {run_context.cur_step}, '
                      f'Time: {run_context.cur_time:.3f}, '
                      f'Total Energy: {self._format_value(te)}, '
                      f'Potential Energy: {self._format_value(pe)}, '
                      f'Kinetic Energy: {self._format_value(ke)}, '
                      f'Temperature: {self._format_value(run_context.temperature)}, ',end=''
                      )
        else:
            print(f'[TorchSPONGE] Step: {run_context.cur_step}, '
                  f'Potential Energy: {self._format_value(run_context.potential)}, ',end=''
                  )
        if run_context.bias is not None:
            print(f'Bias Energy: {self._format_value(run_context.bias)}, ',end='')
        
        if run_context.metrics is not None:
            for name,metric in run_context.metrics.items():
                print(f'{name}: {self._format_value(metric.eval())}, ',end='')
        print(time_str, flush=True)
            
    def _format_value(self, value):
        """
        格式化数值，对标量和 ndarray 保留 4 位小数
        :param value: 输入的数值，可以是标量或 numpy.ndarray
        :return: 格式化后的字符串
        """
        if isinstance(value, np.ndarray):
            if value.size == 1:
                return f"{value.item():.4f}"
            else:
                return "[" + ", ".join([f"{v:.4f}" for v in value.flat]) + "]"
        else:
            return f"{value:.4f}"


class BasicInfo(Callback):
    r"""Callback to print MD simulation information

    Args:
        per_steps (int): How many steps to print information once
        per_epoch (int): How many epochs to print information once
        show_total_time (bool): Whether to show total running time
        show_single_time (bool): Whether to show single step running time
        check_force (bool): Whether to check force

    Supported Platforms:
        ``CPU`` ``GPU``
    """

    def __init__(self,
                 **kwargs
                 ):
        super().__init__()


    def __enter__(self):
        """Return the enter target."""
        return self

    def __exit__(self, *err):
        """Release resources here if have any."""

    def format_info(self,name,param):
        """模块自带的信息打印方法"""
        if param is None:
            return
        for key, value in param.items():
            print(f"[TorchSPONGE] [{name}] {key}: {value}")
        print('---'*30)

    def print_info(self,dt):
        if dt is None:
            return
            
        for k, v in dt.items():
            # 检查是否是module类型且有print_info方法
            if isinstance(v, torch.nn.Module) or isinstance(v, torch.optim.Optimizer):
                if isinstance(v, torch.nn.ModuleList):
                    # 遍历ModuleList中的所有模块
                    for i, module in enumerate(v):
                        if hasattr(module, 'print_info'):
                            module_id = id(module)
                            if module_id not in self.called_modules:
                                # try:
                                info = module.print_info()
                                self.format_info(f"{k}[{i}]", info)
                                self.called_modules.add(module_id)
                                # except Exception as e:
                                #     print(f"[TorchSPONGE] Error printing info for {k}[{i}]: {str(e)}")
                        else:
                            self.print_info(module.__dict__)
                elif hasattr(v, 'print_info'):
                    module_id = id(v)
                    if module_id not in self.called_modules:
                        # try:
                        info = v.print_info()
                        self.format_info(k, info)
                        self.called_modules.add(module_id)
                        # except Exception as e:
                        #     print(f"[TorchSPONGE] Error printing info for {k}: {str(e)}")
                
                # 递归处理子模块
                if hasattr(v, '_modules') and v._modules:
                    self.print_info(v._modules)
                if hasattr(v, '__dict__'):
                    self.print_info(v.__dict__)
                    
            elif isinstance(v, dict):
                self.print_info(v)

    def begin(self, run_context: RunContext):
        """
        Called once before the network executing.

        Args:
            run_context (RunContext): Include some information of the model.
        """
        self.called_modules = set()
        self.print_info(run_context.__dict__)

    def epoch_begin(self, run_context: RunContext):
        """
        Called before each epoch beginning.

        Args:
            run_context (RunContext): Include some information of the model.
        """

    def epoch_end(self, run_context: RunContext):
        """
        Called after each epoch finished.

        Args:
            run_context (RunContext): Include some information of the model.
        """

    def step_begin(self, run_context: RunContext):
        """
        Called before each step beginning.

        Args:
            run_context (RunContext): Include some information of the model.
        """

    def step_end(self, run_context: RunContext):
        """
        Called after each step finished.

        Args:
            run_context (RunContext): Include some information of the model.
        """

    def end(self, run_context: RunContext):
        """
        Called once after network training.

        Args:
            run_context (RunContext): Include some information of the model.
        """

